from dotenv import load_dotenv
from langchain_openai import OpenAIEmbeddings
from pymilvus import MilvusClient

# 直接在代码中定义配置参数
OPENAI_API_KEY = "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"  # 例如："sk-xxxxxxxxxxxxxxxxxxxxxxxx"


# Clash代理设置（根据你的实际端口修改）
# 通常Clash的HTTP代理端口是7890，SOCKS5是7891
PROXY = "http://127.0.0.1:7890"  # 这里使用HTTP代理示例

MILVUS_HOST = "127.0.0.1"
MILVUS_PORT = "19530"
MILVUS_TOKEN = "root:fridge_zzl"  # 如果Milvus没有配置认证，可以设为None


# 定义Milvus操作工具类
class MilvusOperation(object):
    def __init__(self):
        """初始化Milvus客户端和OpenAI Embeddings（使用直接定义的参数）"""
        # 初始化Milvus客户端
        self.client = MilvusClient(
            host=MILVUS_HOST,
            port=MILVUS_PORT,
            token=MILVUS_TOKEN
        )

        # 初始化OpenAI嵌入模型（传入API密钥）
        self.embedding = OpenAIEmbeddings(
            model="text-embedding-3-small",
            api_key=OPENAI_API_KEY  # 直接使用代码中定义的密钥
        )

    def create_table(self):
        """创建Milvus集合（表）"""
        self.client.create_collection(
            collection_name="test",
            dimension=1536,  # 与text-embedding-3-small模型的输出维度匹配
            consistency_level="Strong"
        )


# 主程序
if __name__ == "__main__":
    # 初始化操作类
    milvus_op = MilvusOperation()

    # 创建集合
    milvus_op.create_table()
    print("集合创建成功")

    # 列出集合索引
    res = milvus_op.client.list_indexes(collection_name="test")
    print("集合索引列表：", res)
